#include <utility>
#include "common/common.hpp"
#include "opencv2/opencv.hpp"
#include "model_infer_vino.h"
#include <fstream>

// #include <ipp.h>
// #include <ipp/ippcv.h>
// #include <ipp/ipps.h>

// #include <ipp/ippi.h>
#include <vector>
#include <iostream>

namespace vino_v6
{
    void NmsDetect(
        std::vector<Yolo_VINO::DetRes_VINO> & detections,
        const float nms_threshold_)
    {
        sort(detections.begin(), detections.end(), [](
            const Yolo_VINO::DetRes_VINO &left,
            const Yolo_VINO::DetRes_VINO &right) {
            return left.prob > right.prob;
        });

        for (int i = 0; i < static_cast<int>(detections.size()); i++) {
            for (int j = i + 1; j < static_cast<int>(detections.size()); j++) {
                if (detections[i].classes == detections[j].classes) {
                    float iou = IOUCalculate(detections[i], detections[j]);
                    if (iou > nms_threshold_) {
                        detections[j].prob = 0;
                    }
                }
            }
        }

        detections.erase(std::remove_if(
            detections.begin(),
            detections.end(),
            [](const Yolo_VINO::DetRes_VINO &det) {
            return det.prob == 0;
        }), detections.end());
    }

    float IOUCalculate(
        std::vector<Yolo_VINO::DetRes_VINO>::const_reference det_a,
        std::vector<Yolo_VINO::DetRes_VINO>::const_reference det_b)
    {
        float left = std::max(det_a.x - det_a.w / 2, det_b.x - det_b.w / 2);
        float top = std::max(det_a.y - det_a.h / 2, det_b.y - det_b.h / 2);
        float right = std::min(det_a.x + det_a.w / 2, det_b.x + det_b.w / 2);
        float bottom = std::min(det_a.y + det_a.h / 2, det_b.y + det_b.h / 2);

        float width = std::max(0.0f, right - left);
        float height = std::max(0.0f, bottom - top);

        float inter_area = width * height;

        float area_a = det_a.w * det_a.h;
        float area_b = det_b.w * det_b.h;

        float union_area = area_a + area_b - inter_area;

        return inter_area / union_area;
    }

    std::vector<char> get_file_data(const std::string& file_path) {
        std::ifstream file(file_path, std::ios_base::binary);
        assert(file.good());
        file.seekg(0, std::ios_base::end);
        const int file_len = file.tellg();
        file.seekg(0, std::ios_base::beg);
        std::vector<char> data(file_len, 0);
        file.read(data.data(), file_len);
        file.close();
        return data;
    }


    void decrypt_data(std::vector<char>& data) {
        char key = 0x1;
        for (size_t i = 0; i < data.size(); ++i) {
            data[i] ^= key;
    }
}

    bool decryptFileTo_XB_buffer(
        const std::string& merged_path,
        st_encrypt_info_ovino& info,
        std::vector<char>& xml_data,
        std::vector<char>& bin_data)
    {
        if (merged_path.size() < 6 || merged_path.substr(merged_path.size() - 6) != ".ohdats") {
            throw std::runtime_error("The provided file is not an .hdats file.");
        }

        std::vector<char> merged_data = get_file_data(merged_path);

        if (merged_data.size() < sizeof(st_encrypt_header)) {
            throw std::runtime_error("Merged data is too small to contain the header.");
        }

        st_encrypt_header header{};
        std::memcpy(&header, merged_data.data(), sizeof(st_encrypt_header));

        if (merged_data.size() < sizeof(st_encrypt_header) + header.struct_size + header.xml_size + header.bin_size) {
            throw std::runtime_error("Merged data is too small to contain all parts.");
        }

        std::memcpy(&info, merged_data.data() + sizeof(st_encrypt_header), header.struct_size);

        xml_data.assign(merged_data.begin() + sizeof(st_encrypt_header) + header.struct_size, merged_data.begin() + sizeof(st_encrypt_header) + header.struct_size + header.xml_size);
        decrypt_data(xml_data);

        bin_data.assign(merged_data.begin() + sizeof(st_encrypt_header) + header.struct_size + header.xml_size, merged_data.begin() + sizeof(st_encrypt_header) + header.struct_size + header.xml_size + header.bin_size);
        decrypt_data(bin_data);

        return true;
    }

    ov::Tensor vector_to_tensor(const std::vector<char>& weightsBuffer) {
        size_t weightsSize = weightsBuffer.size();
        ov::Tensor weights(ov::element::u8, {weightsSize});
        std::memcpy(weights.data(), weightsBuffer.data(), weightsSize);
        return weights;
    }

    bool test_inference_with_empty_data(
        ov::InferRequest& infer_request,
        const ov::Shape& input_shape,
        const ov::element::Type& input_type)
    {
        try {
            const ov::Tensor input_tensor(input_type, input_shape);
            const std::vector<float> input_data(
                input_tensor.get_shape().size(), 0.0f);

            std::memcpy(input_tensor.data<float>(), input_data.data(), input_data.size() * sizeof(float));

            infer_request.set_input_tensor(input_tensor);

            const auto start = std::chrono::high_resolution_clock::now();
            infer_request.start_async();
            infer_request.wait_for(std::chrono::milliseconds(1000));
            const auto end = std::chrono::high_resolution_clock::now();

            std::chrono::duration<double, std::milli> latency = end - start;

            return true;
        } catch (const std::exception& e) {
            return false;
        }
    }

	Yolo_VINO::Yolo_VINO() = default;

	Yolo_VINO::~Yolo_VINO() = default;

	void Yolo_VINO::Reset() {
		for (auto &req : infer_requests) {
			req = {};
		}
		compiled_model = {};
	}

    bool Yolo_VINO::Init(
        const char* hdats_path)
    {
        core = Core();

        const std::string file_path = hdats_path;

        std::vector<char> xml_buffer = {};
        std::vector<char> bin_buffer = {};

        decryptFileTo_XB_buffer(
            hdats_path,
            o_m_info,
            xml_buffer,
            bin_buffer);

        std::vector<std::pair<int, int>> grids;

        for (int i = 0; i < o_m_info.stride_num; ++i) {
            std::vector<int> grid_size(2);
            grid_size[0] = o_m_info.image_height / o_m_info.strides[i];
            grid_size[1] = o_m_info.image_width / o_m_info.strides[i];
            m_grids.push_back(grid_size);
        }

        std::cout << "Loading model ok" << std::endl;

        std::cout << "Parsed XML Data Size: " << xml_buffer.size() << std::endl;
        std::cout << "Parsed BIN Data Size: " << bin_buffer.size() << std::endl;

        std::string model_xml(xml_buffer.begin(), xml_buffer.end());
        ov::Tensor bin_tensor = vector_to_tensor(bin_buffer);

        try {
            std::shared_ptr<ov::Model> model = core.read_model(model_xml, bin_tensor);
            compiled_model = core.compile_model(model, "AUTO");
        }
        catch (const std::exception& e) {
            throw;
        }

        size_t num_requests = o_m_info.batch_size;
        infer_requests.resize(num_requests);
        for (size_t i = 0; i < num_requests; ++i) {
            infer_requests[i] = compiled_model.create_infer_request();
        }
        if (!test_inference_with_empty_data(
            infer_requests[0],
            compiled_model.input().get_shape(),
            compiled_model.input().get_element_type())) {
            return false;
        }
        return true;
    }

	int Yolo_VINO::BatchSize() const {
		return o_m_info.batch_size;
	}


    void Yolo_VINO::Detect(
        st_dlm_data* p_imgs,
        int num,
        std::vector<std::vector<DetRes_VINO>>& vv_results)
    {
        std::vector<cv::Mat> mats(num);

        for (int i = 0; i < num; ++i) {
            const st_dlm_data& img_data = p_imgs[i];
            mats[i] = img_data.img.rgb8;
        }

        int net_w = o_m_info.image_width,
            net_h = o_m_info.image_height,
            net_c = o_m_info.input_channel;

        cv::Scalar mean_val(0, 0, 0);
        if (net_c == 1) {
            mean_val = cv::Scalar(0);
        }

        std::vector<cv::Mat> prepared_mats = preprocess_image(
            mats,
            net_w,
            net_h,
            mean_val);

        std::vector<ov::Tensor> V_Tensor;
        vv_results.clear();

        for (size_t i = 0; i < prepared_mats.size(); ++i) {
            Output<const ov::Node> input_port = compiled_model.input();
            ov::Tensor input_tensor(
                input_port.get_element_type(),
                input_port.get_shape(),
                prepared_mats[i].ptr(0));

            infer_requests[0].set_input_tensor(input_tensor);

            infer_requests[0].infer();

            ov::Tensor output = infer_requests[0].get_output_tensor(0);

            V_Tensor.push_back(output);

            std::vector<DetRes_VINO> detection_results = postProcess_yolov6(
                mats[i],
                output,
                o_m_info.obj_threshold,
                o_m_info.nms_threshold);
            vv_results.push_back(detection_results);
        }

    }

    std::vector<std::vector<Yolo_VINO::DetRes_VINO>> Yolo_VINO::postProcess_batch(
        const std::vector<cv::Mat> &img_batch,
        const std::vector<ov::Tensor> &data_tensors,
        const float prob_threshold_,
        const float nms_threshold_)
    {
        std::vector<std::vector<DetRes_VINO>> batch_results;

        assert(img_batch.size() == data_tensors.size());

        for (size_t i = 0; i < img_batch.size(); ++i) {
            const cv::Mat& src_img = img_batch[i];
            const ov::Tensor& data_tensor = data_tensors[i];

            std::vector<DetRes_VINO> results =
                postProcess_yolov6(
                    src_img,
                    data_tensor,
                    prob_threshold_,
                    nms_threshold_);

            batch_results.push_back(results);
        }

        return batch_results;
    }

    std::vector<Yolo_VINO::DetRes_VINO> Yolo_VINO::postProcess_yolov6(
        const cv::Mat &src_img,
        const ov::Tensor& data_tensor,
        const float prob_threshold_,
        const float nms_threshold_)
    {
        auto* output = data_tensor.data<float>();

        std::vector<DetRes_VINO> result;
        const float ratio =
            std::max(
                static_cast<float>(src_img.cols) / static_cast<float>(o_m_info.image_width),
                static_cast<float>(src_img.rows) / static_cast<float>(o_m_info.image_height));
        const float scale = std::min(
            static_cast<float>(o_m_info.image_width) / static_cast<float>(src_img.cols),
            static_cast<float>(o_m_info.image_height) / static_cast<float>(src_img.rows));

        float padd_h = static_cast<float>(
            std::round((static_cast<float>(o_m_info.image_height) -
                static_cast<float>(src_img.rows) * scale) / 2));
        float padd_w = static_cast<float>(
            std::round((static_cast<float>(o_m_info.image_width) -
                static_cast<float>(src_img.cols) * scale) / 2));

        for (auto& grid : m_grids) {
            for (int h = 0; h < grid[0]; h++) {
                for (int w = 0; w < grid[1]; w++) {
                    float* row = output + ((h * grid[1] + w) * (CATEGORY + 5));
                    DetRes_VINO box{};
                    auto max_pos = std::max_element(row + 5, row + CATEGORY + 5);
                    box.prob = row[4] * row[max_pos - row];
                    if (box.prob < prob_threshold_) continue;
                    box.classes = static_cast<int>(max_pos - row - 5);
                    box.x = (row[0] - padd_w) * ratio;
                    box.y = (row[1] - padd_h) * ratio;
                    box.w = row[2] * ratio;
                    box.h = row[3] * ratio;
                    result.push_back(box);
                }
            }
        }

        NmsDetect(result, nms_threshold_);
        return result;
    }

    void visualize_and_save_yolov6(
        const cv::Mat& original_img,
        const std::vector<Yolo_VINO::DetRes_VINO>& detection_results,
        const std::string& output_file_path)
    {
        std::vector<cv::Scalar> class_colors;

        class_colors.resize(1);
        srand(static_cast<unsigned int>(time(nullptr)));
        for (cv::Scalar& class_color : class_colors) {
            class_color = cv::Scalar(rand() % 255, rand() % 255, rand() % 255);
        }

        if (original_img.empty()) {
            std::cerr << "Error: Original image is empty." << std::endl;
            return;
        }

        cv::Mat img = original_img.clone();

        for (const auto& [classes, x, y, w, h, prob] : detection_results) {
            cv::Rect box(x - w / 2, y - h / 2, w, h);

            if (const int class_id = classes; class_id >= 0 && class_id < static_cast<int>(1)) {
                cv::rectangle(img, box, class_colors[class_id % class_colors.size()], 2, 8);

                std::string label = "1:" + std::to_string(prob).substr(0, 4);

                const cv::Size textSize = cv::getTextSize(label, cv::FONT_HERSHEY_SIMPLEX, 0.5, 1, 0);

                const cv::Rect textBox(box.tl().x, box.tl().y - 15, textSize.width, textSize.height + 5);

                cv::rectangle(img, textBox, class_colors[class_id % class_colors.size()], cv::FILLED);

                cv::putText(
                    img, label,
                    cv::Point(box.tl().x, box.tl().y - 5),
                    cv::FONT_HERSHEY_SIMPLEX, 0.5,
                    cv::Scalar(255, 255, 255));
            } else {
                std::cerr << "Error: Class ID out of range." << std::endl;
            }
        }

        if (!cv::imwrite(output_file_path, img)) {
            std::cerr << "Error: Failed to save image to " << output_file_path << std::endl;
        } else {
            std::cout << "Detection results have been saved to: " << output_file_path << std::endl;
        }
    }

}//namespace yolo_v6